#include "stdafx.h"
#include "Utils.h"
#include "Holders.h" 

void WMReaderDllUtils::SplitKeyPath(const std::wstring &keyPath,
                                          std::wstring &preDefKey,
                                          std::wstring &subKey)
{
    size_t len = keyPath.length();
    size_t index = keyPath.find_first_of(L"\\");

    if(index != std::string::npos) {
        preDefKey = keyPath.substr(0, index);
        subKey = keyPath.substr(index + 1, len - 1);
    } else {
        preDefKey = keyPath;
        subKey = L"";
    }                 
}

#ifdef UNDER_CE

void WMReaderDllUtils::PreDefKeyFromString(HKEY &preDefKey, const std::wstring &preDefString)
{
    if(preDefString == L"HKEY_LOCAL_MACHINE") {
        preDefKey = HKEY_LOCAL_MACHINE;
    } else if(preDefString == L"HKEY_CURRENT_USER") {
        preDefKey = HKEY_CURRENT_USER;
    } else if(preDefString == L"HKEY_CLASSES_ROOT") {
        preDefKey = HKEY_CLASSES_ROOT;
    } else if(preDefString == L"HKEY_USERS") {
        preDefKey = HKEY_USERS; 
    }
}

void WMReaderDllUtils::GetFileSystem(BYTE *input, IRAPIStream *stream)
{
    if(stream != NULL) {
        std::auto_ptr <RapiStreamHolder>
        rapiStream(new RapiStreamHolder(stream));

        WIN32_FIND_DATA fileSystemItem;
        memset(&fileSystemItem, 0, sizeof(fileSystemItem));

        DWORD len = sizeof(fileSystemItem);
        DWORD bytes = 0;

        std::auto_ptr <FindHolder> findHolder(new FindHolder);
        std::wstring path = reinterpret_cast<const wchar_t*>(input);

        HRESULT result = S_OK;

        result = findHolder->FirstFile(path, fileSystemItem);
        if(result == S_OK) {
            do {
                rapiStream->Stream()->Write(&len, sizeof(len), &bytes);
                rapiStream->Stream()->Write(&fileSystemItem, len, &bytes);

                memset(&fileSystemItem, 0, sizeof(fileSystemItem));
                bytes = 0;
            } while(findHolder->NextFile(fileSystemItem) == S_OK);
        }
        len = connect_set::lastBuffer;
        rapiStream->Stream()->Write(&len, sizeof(len), &bytes);
    }
}

void WMReaderDllUtils::GetRegistryKeysList(BYTE* input, IRAPIStream *stream)
{
    if(stream != NULL) {
        std::auto_ptr <RapiStreamHolder>
        rapiStream(new RapiStreamHolder(stream));

        std::wstring keyPath, preDefKey, subKey;

        if(input != NULL) {
            keyPath = reinterpret_cast<const wchar_t*>(input);
            WMReaderDllUtils::SplitKeyPath(keyPath, preDefKey, subKey);
        }                

        HKEY key;
        WMReaderDllUtils::PreDefKeyFromString(key, preDefKey);

        DWORD bytes = 0;
        DWORD len = 0;
        DWORD childLen = MAX_PATH;

        registry_set::REGKEY childKey;
        memset(&childKey, 0, sizeof(childKey));

        std::auto_ptr <RegKeyHolder> keyHolder(new RegKeyHolder);

        if(keyHolder->OpenKey(key, subKey) == S_OK) {

            for (size_t i = 0; RegEnumKeyEx(keyHolder->Key(), i, childKey.regKey,
                &childLen, NULL, NULL, NULL, NULL) == ERROR_SUCCESS; ++i) {

                len = childLen * sizeof(wchar_t);
                rapiStream->Stream()->Write(&len, sizeof(len), &bytes);
                rapiStream->Stream()->Write(&childKey, len, &bytes);

                memset(&childKey, 0, sizeof(childKey));
                childLen = MAX_PATH;
                bytes = 0;
            }
        }
        len = connect_set::lastBuffer;
        rapiStream->Stream()->Write(&len, sizeof(len), &bytes);
    } 
}

void WMReaderDllUtils::GetRegistryKeyValues(BYTE *input, IRAPIStream *stream)
{
    if(stream != NULL) {
        std::auto_ptr <RapiStreamHolder>
        rapiStream(new RapiStreamHolder(stream));

        std::wstring keyPath, preDefKey, subKey;

        if(input != NULL) {
            keyPath = reinterpret_cast<const wchar_t*>(input);
            WMReaderDllUtils::SplitKeyPath(keyPath, preDefKey, subKey);
        }                

        HKEY key;
        WMReaderDllUtils::PreDefKeyFromString(key, preDefKey);

        registry_set::REGKEYVALUE keyValue;
        memset(&keyValue, 0, sizeof(keyValue));

        DWORD len = sizeof(keyValue);
        DWORD bytes = 0;
        DWORD valueSize = MAX_PATH;

        std::auto_ptr <RegKeyHolder> keyHolder(new RegKeyHolder);

        if(keyHolder->OpenKey(key, subKey) == S_OK) {

            keyValue.bufferSize = registry_set::maxValueBufferSize;
            for(size_t i = 0; RegEnumValue(keyHolder->Key(), i, keyValue.valueName, &valueSize, NULL,
                &keyValue.valueType, keyValue.valueBuffer, &keyValue.bufferSize) == ERROR_SUCCESS; ++i) {

                rapiStream->Stream()->Write(&len, sizeof(len), &bytes);
                rapiStream->Stream()->Write(&keyValue, len, &bytes);

                memset(&keyValue, 0, sizeof(keyValue));
                keyValue.bufferSize = registry_set::maxValueBufferSize;
                valueSize = MAX_PATH;
                bytes = 0;            
            }
        }
        len = connect_set::lastBuffer;
        rapiStream->Stream()->Write(&len, sizeof(len), &bytes); 
    }
}

void WMReaderDllUtils::GetDeviceInfo(IRAPIStream *stream)
{
    if(stream != NULL) {
        std::auto_ptr <RapiStreamHolder>
        rapiStream(new RapiStreamHolder(stream));

        device_info::DEVICEINFO deviceInfo;
        memset(&deviceInfo, 0, sizeof(deviceInfo));

        OSVERSIONINFO osVersion;
        memset(&osVersion, 0, sizeof(osVersion));

        osVersion.dwOSVersionInfoSize = sizeof(OSVERSIONINFO);
        GetVersionEx(&osVersion);

        deviceInfo.osMinVersion = osVersion.dwMinorVersion;
        deviceInfo.osMajorVersion = osVersion.dwMajorVersion;
        deviceInfo.osBuildNumber = osVersion.dwBuildNumber;    

        SYSTEM_INFO sysInfo;
        memset(&sysInfo, 0, sizeof(sysInfo));

        GetSystemInfo(&sysInfo);
        deviceInfo.procType = sysInfo.dwProcessorType;

        MEMORYSTATUS memStat;
        memset(&memStat, 0, sizeof(memStat));

        memStat.dwLength = sizeof(MEMORYSTATUS);
        GlobalMemoryStatus(&memStat);

        deviceInfo.totalPhysMem = memStat.dwTotalPhys;
        deviceInfo.availPhysMem = memStat.dwAvailPhys;
        deviceInfo.totalVirtMem = memStat.dwTotalVirtual;
        deviceInfo.availVirtMem = memStat.dwAvailVirtual;
        deviceInfo.memoryLoad = memStat.dwMemoryLoad;

        DEVMODE devMode;
        memset(&devMode, 0, sizeof(devMode));

        devMode.dmSize = sizeof(DEVMODE);
        EnumDisplaySettings(NULL, ENUM_CURRENT_SETTINGS, &devMode);

        deviceInfo.displayWidth = devMode.dmPelsWidth;
        deviceInfo.displayHeight = devMode.dmPelsHeight;
        deviceInfo.displayBits = devMode.dmBitsPerPel;

        SYSTEM_POWER_STATUS_EX batStatus;
        memset(&batStatus, 0, sizeof(SYSTEM_POWER_STATUS_EX));

        GetSystemPowerStatusEx(&batStatus, TRUE);
        deviceInfo.batStatus = batStatus.BatteryLifePercent;

        DWORD len  = sizeof(device_info::DEVICEINFO);
        DWORD bytes = 0;

        rapiStream->Stream()->Write(&len, sizeof(len), &bytes);
        rapiStream->Stream()->Write(&deviceInfo, len, &bytes);

        len = connect_set::lastBuffer;
        rapiStream->Stream()->Write(&len, sizeof(len), &bytes);
    }
}

#endif